#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/8/30 7:05
@Author   : liujingmao
@File     : 1.多LLM链选择示例.py
"""

import dotenv
from langchain_community.chat_models import QianfanChatEndpoint
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import ConfigurableField
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

# 1. 创建提示模板 定义默认大语言模型
prompt = ChatPromptTemplate.from_template("{query}")

# 2. 构建LLM
llm = ChatOpenAI(model="gpt-3.5-turbo-16k").configurable_alternatives(
    ConfigurableField(id="llm"),
    default_key="gpt-3.5",
    gpt4=ChatOpenAI(model="gpt-4o"),
    gpt35=ChatOpenAI(model="gpt-3.5-turbo"),
    wenxin=QianfanChatEndpoint(),
)

# 3.构建链应用
chain = prompt | llm | StrOutputParser()

# 4.调用链并传递配置信息，并切换到文心一言模型或者gpt4模型
content = chain.invoke(
    {"query": "你好，请简单介绍一下你自己，你是哪家公司开发的，你们公司的CEO是雷军吗？"},
    config={"configurable": {"llm": "wenxin"}}
)
print(content)
